#define STRICT
#include "handling.h"
#include "hooked.h"
#include <psapi.h>

HMODULE injectModule;

void Patch(DWORD* address, DWORD data)
{
	DWORD oldProtect;
	VirtualProtect(address, sizeof(DWORD), PAGE_READWRITE, &oldProtect);
	*address = data;
	VirtualProtect(address, sizeof(DWORD), oldProtect, &oldProtect);
}

BOOL CALLBACK DllMain(HINSTANCE instance, DWORD reason, VOID* reserved)
{
	switch(reason)
	{
	case DLL_PROCESS_ATTACH:
		injectModule = instance;
		break;
	}
	return TRUE;
}

void HookModuleExport(char* module, PIMAGE_EXPORT_DIRECTORY directory)
{
	size_t functionsCount = directory->NumberOfFunctions;
	DWORD* functions = (DWORD*)(module + directory->AddressOfFunctions);
	for(size_t i = 0; i < functionsCount; ++i)
	{
		HookedFunction* function = GetHookedFunction((void*)(module + functions[i]));
		if(!function) continue;
		Patch(functions + i, (DWORD)function->function - (DWORD)module);
	}
}

void HookModuleImport(char* module, PIMAGE_IMPORT_DESCRIPTOR descriptor)
{
	for(; descriptor->Name; ++descriptor)
	{
		DWORD* iat = (DWORD*)(module + descriptor->FirstThunk);

		for(unsigned i = 0; iat[i]; ++i)
		{
			HookedFunction* function = GetHookedFunction((void*)iat[i]);
			if(!function) continue;
			Patch(iat + i, (DWORD)function->function);
		}
	}
}

void InjectToModule(char* module)
{
	if(module == (char*)injectModule) return;

	PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)module;
	PIMAGE_FILE_HEADER fileHeader = (PIMAGE_FILE_HEADER)(module + dosHeader->e_lfanew + 4);
	PIMAGE_OPTIONAL_HEADER optionalHeader = (PIMAGE_OPTIONAL_HEADER)(fileHeader + 1);

	IMAGE_DATA_DIRECTORY& exportDirectory = optionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
	if(exportDirectory.VirtualAddress && exportDirectory.Size)
		HookModuleExport(module, (PIMAGE_EXPORT_DIRECTORY)(module + exportDirectory.VirtualAddress));

	IMAGE_DATA_DIRECTORY& importDirectory = optionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
	if(importDirectory.VirtualAddress && importDirectory.Size)
		HookModuleImport(module, (PIMAGE_IMPORT_DESCRIPTOR)(module + importDirectory.VirtualAddress));
}

void InjectToAllModules()
{
	static HMODULE modules[1024];
	DWORD needed;
	EnumProcessModules(GetCurrentProcess(), modules, sizeof(modules), &needed);

	unsigned modulesCount = min(needed, sizeof(modules)) / sizeof(HMODULE);

	for(unsigned i = 0; i < modulesCount; ++i)
		InjectToModule((char*)modules[i]);
}

size_t GetProgramEntryPoint()
{
	char* module = (char*)GetModuleHandle(NULL);
	PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)module;
	PIMAGE_FILE_HEADER fileHeader = (PIMAGE_FILE_HEADER)(module + dosHeader->e_lfanew + 4);
	PIMAGE_OPTIONAL_HEADER optionalHeader = (PIMAGE_OPTIONAL_HEADER)(fileHeader + 1);

	return (DWORD)(module + optionalHeader->AddressOfEntryPoint);
}
